#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
import re
import csv
from datetime import datetime
from scapy.all import *
from scapy.layers.inet import IP, TCP, UDP, ICMP
from scapy.layers.inet6 import IPv6
from scapy.layers.l2 import Ether, ARP
from scapy.layers.http import HTTP, HTTPRequest, HTTPResponse
from scapy.layers.dns import DNS

class PacketParser:
    """
    数据包解析类，负责解析捕获的数据包
    """
    
    def __init__(self):
        """
        初始化数据包解析器
        """
        # 协议映射字典，用于存储已知的协议
        self.protocol_map = {
            "Ether": "以太网",
            "IP": "IPv4",
            "IPv6": "IPv6",
            "TCP": "TCP",
            "UDP": "UDP",
            "ICMP": "ICMP",
            "ICMPv6": "ICMPv6",
            "HTTP": "HTTP",
            "HTTPRequest": "HTTP请求",
            "HTTPResponse": "HTTP响应",
            "DNS": "DNS",
            "DHCP": "DHCP",
            "ARP": "ARP",
            "Raw": "数据",
            "Padding": "填充",
            "LLC": "LLC",
            "SNAP": "SNAP",
            "BOOTP": "BOOTP"
        }
        
        # 常用端口映射
        self.port_map = {
            20: "FTP数据",
            21: "FTP控制",
            22: "SSH",
            23: "Telnet",
            25: "SMTP",
            53: "DNS",
            67: "DHCP服务器",
            68: "DHCP客户端",
            69: "TFTP",
            80: "HTTP",
            110: "POP3",
            123: "NTP",
            143: "IMAP",
            161: "SNMP",
            162: "SNMP陷阱",
            443: "HTTPS",
            465: "SMTPS",
            993: "IMAPS",
            995: "POP3S",
            3389: "RDP"
        }
    
    def parse_packet(self, packet):
        """
        解析单个数据包，提取关键信息
        
        Args:
            packet: scapy捕获的数据包对象
            
        Returns:
            dict: 包含解析后信息的字典
        """
        if packet is None:
            return {
                "error": "无效的数据包（空值）",
                "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                "protocol": "错误",
                "length": 0,
                "raw_packet": None
            }
            
        try:
            # 基本信息
            parsed_data = {
                "number": getattr(packet, "number", "N/A"),
                "time": self._format_timestamp(packet),
                "length": len(packet),
                "layers": [],
                "protocols": [],
                "raw_packet": packet  # 保存原始数据包以便详细查看
            }
            
            # 提取协议层和协议类型
            layers = []
            current_layer = packet
            max_layers = 10  # 防止无限循环
            layer_count = 0
            
            while current_layer and layer_count < max_layers:
                layer_count += 1
                try:
                    layer_name = current_layer.name
                    if layer_name not in parsed_data["protocols"]:
                        parsed_data["protocols"].append(self._get_protocol_name(layer_name))
                    
                    # 解析当前层
                    layer_info = self._parse_layer(current_layer)
                    if layer_info:
                        parsed_data["layers"].append(layer_info)
                    
                    # 检查是否有更多层
                    if hasattr(current_layer, "payload") and current_layer.payload:
                        current_layer = current_layer.payload
                    else:
                        break
                except Exception as layer_error:
                    # 记录层解析错误但继续处理
                    parsed_data["layers"].append({
                        "name": "解析错误",
                        "error": f"解析层错误: {str(layer_error)}"
                    })
                    break
            
            # 如果达到最大层数，添加警告
            if layer_count >= max_layers:
                parsed_data["warning"] = "数据包层数过多，可能存在循环引用"
            
            # 提取源和目的地址信息
            try:
                self._extract_src_dst(packet, parsed_data)
            except Exception as extract_error:
                parsed_data["src"] = "未知"
                parsed_data["dst"] = "未知"
                parsed_data["error_extract"] = f"提取地址错误: {str(extract_error)}"
            
            # 确定顶层协议（应用层）
            if parsed_data["protocols"]:
                parsed_data["protocol"] = parsed_data["protocols"][-1]
            else:
                parsed_data["protocol"] = "未知"
            
            return parsed_data
            
        except Exception as e:
            # 确保即使在出现严重错误的情况下也返回一个有效的字典
            try:
                length = len(packet) if packet else 0
            except:
                length = 0
                
            try:
                timestamp = self._format_timestamp(packet)
            except:
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
                
            return {
                "error": f"解析错误: {str(e)}",
                "time": timestamp,
                "protocol": "错误",
                "length": length,
                "raw_packet": packet,
                "src": "未知",
                "dst": "未知"
            }
    
    def _format_timestamp(self, packet):
        """
        格式化数据包的时间戳
        """
        try:
            # scapy数据包时间戳在packet.time
            if hasattr(packet, "time"):
                timestamp = float(packet.time)
                return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        except:
            pass
        
        # 默认返回当前时间
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    
    def _get_protocol_name(self, protocol_key):
        """
        根据协议键获取可读的协议名称
        """
        return self.protocol_map.get(protocol_key, protocol_key)
    
    def _parse_layer(self, layer):
        """
        解析单个协议层
        
        Args:
            layer: 协议层对象
            
        Returns:
            dict: 包含层信息的字典
        """
        try:
            layer_name = layer.name
            layer_info = {
                "name": self._get_protocol_name(layer_name),
                "fields": {}
            }
            
            # 获取该层的字段 - 使用scapy的fields属性
            if hasattr(layer, "fields"):
                for field_name, field_value in layer.fields.items():
                    # 将值转换为字符串
                    if isinstance(field_value, bytes):
                        # 尝试解码字节为字符串
                        try:
                            field_value = field_value.decode('utf-8', errors='replace')
                        except:
                            field_value = str(field_value)
                    
                    layer_info["fields"][field_name] = str(field_value)
            
            # 根据不同的协议层类型进行特殊处理
            if isinstance(layer, Ether):
                self._enhance_ether_layer(layer, layer_info)
            elif isinstance(layer, IP):
                self._enhance_ip_layer(layer, layer_info)
            elif isinstance(layer, IPv6):
                self._enhance_ipv6_layer(layer, layer_info)
            elif isinstance(layer, TCP):
                self._enhance_tcp_layer(layer, layer_info)
            elif isinstance(layer, UDP):
                self._enhance_udp_layer(layer, layer_info)
            elif isinstance(layer, ICMP):
                self._enhance_icmp_layer(layer, layer_info)
            elif isinstance(layer, ARP):
                self._enhance_arp_layer(layer, layer_info)
            elif isinstance(layer, DNS):
                self._enhance_dns_layer(layer, layer_info)
            elif layer_name == "Raw":
                self._enhance_raw_layer(layer, layer_info)
            
            return layer_info
            
        except Exception as e:
            return {
                "name": "未知层",
                "error": f"解析层错误: {str(e)}"
            }
    
    def _enhance_ether_layer(self, layer, layer_info):
        """
        增强以太网层信息
        """
        try:
            layer_info["source_mac"] = layer.src
            layer_info["dest_mac"] = layer.dst
            layer_info["type"] = hex(layer.type)
        except:
            pass
    
    def _enhance_ip_layer(self, layer, layer_info):
        """
        增强IP层信息
        """
        try:
            layer_info["source_ip"] = layer.src
            layer_info["dest_ip"] = layer.dst
            layer_info["ttl"] = layer.ttl
            layer_info["version"] = layer.version
            layer_info["ihl"] = layer.ihl
            layer_info["tos"] = layer.tos
            layer_info["id"] = layer.id
            layer_info["flags"] = layer.flags
            layer_info["frag"] = layer.frag
            layer_info["proto"] = layer.proto
        except:
            pass
    
    def _enhance_ipv6_layer(self, layer, layer_info):
        """
        增强IPv6层信息
        """
        try:
            layer_info["source_ip"] = layer.src
            layer_info["dest_ip"] = layer.dst
            layer_info["version"] = layer.version
            layer_info["tc"] = layer.tc
            layer_info["fl"] = layer.fl
            layer_info["nh"] = layer.nh
            layer_info["hlim"] = layer.hlim
        except:
            pass
    
    def _enhance_tcp_layer(self, layer, layer_info):
        """
        增强TCP层信息
        """
        try:
            # 添加TCP层的关键信息
            src_port = layer.sport
            dst_port = layer.dport
            
            layer_info["source_port"] = src_port
            layer_info["dest_port"] = dst_port
            
            if src_port in self.port_map:
                layer_info["source_service"] = self.port_map[src_port]
            
            if dst_port in self.port_map:
                layer_info["dest_service"] = self.port_map[dst_port]
                
            layer_info["seq"] = layer.seq
            layer_info["ack"] = layer.ack
                
            # 解析TCP标志
            flags = []
            if layer.flags:
                flags_value = layer.flags
                if 'S' in flags_value:
                    flags.append("SYN")
                if 'A' in flags_value:
                    flags.append("ACK")
                if 'F' in flags_value:
                    flags.append("FIN")
                if 'R' in flags_value:
                    flags.append("RST")
                if 'P' in flags_value:
                    flags.append("PSH")
                if 'U' in flags_value:
                    flags.append("URG")
                
                layer_info["flags"] = flags
        except:
            pass
    
    def _enhance_udp_layer(self, layer, layer_info):
        """
        增强UDP层信息
        """
        try:
            # 添加UDP层的关键信息
            src_port = layer.sport
            dst_port = layer.dport
            
            layer_info["source_port"] = src_port
            layer_info["dest_port"] = dst_port
            
            if src_port in self.port_map:
                layer_info["source_service"] = self.port_map[src_port]
            
            if dst_port in self.port_map:
                layer_info["dest_service"] = self.port_map[dst_port]
                
            layer_info["len"] = layer.len
            layer_info["chksum"] = layer.chksum
        except:
            pass
    
    def _enhance_icmp_layer(self, layer, layer_info):
        """
        增强ICMP层信息
        """
        try:
            layer_info["type"] = layer.type
            layer_info["code"] = layer.code
            
            # 解释ICMP类型
            icmp_types = {
                0: "Echo Reply",
                3: "Destination Unreachable",
                5: "Redirect",
                8: "Echo Request",
                11: "Time Exceeded"
            }
            
            if layer.type in icmp_types:
                layer_info["type_desc"] = icmp_types[layer.type]
        except:
            pass
    
    def _enhance_arp_layer(self, layer, layer_info):
        """
        增强ARP层信息
        """
        try:
            layer_info["hw_type"] = layer.hwtype
            layer_info["proto_type"] = layer.ptype
            layer_info["hw_len"] = layer.hwlen
            layer_info["proto_len"] = layer.plen
            layer_info["op"] = layer.op
            
            layer_info["sender_mac"] = layer.hwsrc
            layer_info["sender_ip"] = layer.psrc
            layer_info["target_mac"] = layer.hwdst
            layer_info["target_ip"] = layer.pdst
            
            # 解释ARP操作
            arp_ops = {
                1: "who-has (请求)",
                2: "is-at (响应)",
                3: "RARP请求",
                4: "RARP响应"
            }
            
            if layer.op in arp_ops:
                layer_info["op_desc"] = arp_ops[layer.op]
        except:
            pass
    
    def _enhance_dns_layer(self, layer, layer_info):
        """
        增强DNS层信息
        """
        try:
            layer_info["id"] = layer.id
            layer_info["qr"] = layer.qr
            
            # 是查询还是响应
            if layer.qr == 0:
                layer_info["type"] = "Query"
            else:
                layer_info["type"] = "Response"
                
            # 查询类型
            if hasattr(layer, "qd") and layer.qd:
                queries = []
                for i in range(layer.qdcount):
                    if i < len(layer.qd):
                        query = layer.qd[i]
                        queries.append({
                            "name": query.qname.decode('utf-8', errors='replace') if isinstance(query.qname, bytes) else str(query.qname),
                            "type": query.qtype
                        })
                if queries:
                    layer_info["queries"] = queries
                    
            # 响应内容
            if hasattr(layer, "an") and layer.an:
                answers = []
                for i in range(layer.ancount):
                    if i < len(layer.an):
                        answer = layer.an[i]
                        answers.append({
                            "name": answer.rrname.decode('utf-8', errors='replace') if isinstance(answer.rrname, bytes) else str(answer.rrname),
                            "type": answer.type,
                            "data": str(answer.rdata)
                        })
                if answers:
                    layer_info["answers"] = answers
        except:
            pass
    
    def _enhance_raw_layer(self, layer, layer_info):
        """
        增强Raw层信息
        """
        try:
            # 尝试解析HTTP等应用层协议
            payload = layer.load
            if isinstance(payload, bytes):
                # 尝试解码为UTF-8文本
                try:
                    text = payload.decode('utf-8', errors='replace')
                    layer_info["data"] = text[:200] + "..." if len(text) > 200 else text
                    
                    # 尝试识别HTTP
                    if text.startswith("GET ") or text.startswith("POST ") or text.startswith("HTTP/"):
                        layer_info["likely_protocol"] = "HTTP"
                        
                        # 提取HTTP头部
                        headers = {}
                        for line in text.split("\r\n"):
                            if not line or line.startswith("GET ") or line.startswith("POST "):
                                continue
                            if ":" in line:
                                key, value = line.split(":", 1)
                                headers[key.strip()] = value.strip()
                        
                        if headers:
                            layer_info["http_headers"] = headers
                except:
                    # 如果不是文本，就保留原始字节形式
                    layer_info["data"] = str(payload)
        except:
            pass
    
    def _extract_src_dst(self, packet, parsed_data):
        """
        从数据包中提取源和目的地址信息
        
        Args:
            packet: 数据包对象
            parsed_data: 解析数据字典，将添加src和dst字段
        """
        # 默认值
        src = "未知"
        dst = "未知"
        src_port = ""
        dst_port = ""
        
        # 尝试提取MAC地址
        if Ether in packet:
            parsed_data["src_mac"] = packet[Ether].src
            parsed_data["dst_mac"] = packet[Ether].dst
        
        # 尝试提取IP地址
        if IP in packet:
            src = packet[IP].src
            dst = packet[IP].dst
        elif IPv6 in packet:
            src = packet[IPv6].src
            dst = packet[IPv6].dst
        elif ARP in packet:
            src = packet[ARP].psrc
            dst = packet[ARP].pdst
        
        # 尝试提取端口信息
        if TCP in packet:
            src_port = str(packet[TCP].sport)
            dst_port = str(packet[TCP].dport)
        elif UDP in packet:
            src_port = str(packet[UDP].sport)
            dst_port = str(packet[UDP].dport)
        
        # 添加端口信息（如果有）
        if src_port:
            parsed_data["src"] = f"{src}:{src_port}"
            parsed_data["src_port"] = src_port
        else:
            parsed_data["src"] = src
            
        if dst_port:
            parsed_data["dst"] = f"{dst}:{dst_port}"
            parsed_data["dst_port"] = dst_port
        else:
            parsed_data["dst"] = dst
        
        # 添加IP地址（如果有）
        if src != "未知":
            parsed_data["src_ip"] = src
        if dst != "未知":
            parsed_data["dst_ip"] = dst
    
    def export_to_json(self, packets, filename):
        """
        将数据包导出为JSON格式
        
        Args:
            packets: 数据包列表
            filename: 输出文件名
            
        Returns:
            bool: 是否成功导出
        """
        try:
            export_data = []
            
            for packet in packets:
                # 创建一个可序列化的副本
                if isinstance(packet, dict):
                    packet_copy = {}
                    
                    # 复制除了raw_packet之外的所有数据
                    for key, value in packet.items():
                        if key != 'raw_packet':
                            # 确保所有数据都可以JSON序列化
                            if isinstance(value, bytes):
                                try:
                                    packet_copy[key] = value.decode('utf-8', errors='replace')
                                except:
                                    packet_copy[key] = str(value)
                            elif isinstance(value, (dict, list, str, int, float, bool, type(None))):
                                packet_copy[key] = value
                            else:
                                packet_copy[key] = str(value)
                    
                    export_data.append(packet_copy)
            
            # 使用UTF-8-SIG编码写入，添加BOM标记以确保Windows正确识别
            with open(filename, 'w', encoding='utf-8-sig') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=2)
                
            return True
        except Exception as e:
            print(f"导出JSON失败: {str(e)}")
            return False
    
    def export_to_csv(self, packets, filename):
        """
        将数据包导出为CSV格式
        
        Args:
            packets: 数据包列表
            filename: 输出文件名
            
        Returns:
            bool: 是否成功导出
        """
        try:
            # 定义CSV表头
            fieldnames = ['number', 'time', 'protocol', 'src', 'dst', 'length']
            
            # 使用UTF-8-SIG编码写入，添加BOM标记以确保Windows正确识别
            with open(filename, 'w', newline='', encoding='utf-8-sig') as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                
                for packet in packets:
                    # 创建一个只包含表头字段的字典
                    row = {}
                    for field in fieldnames:
                        value = packet.get(field, '')
                        # 确保所有数据都是字符串
                        if isinstance(value, bytes):
                            try:
                                row[field] = value.decode('utf-8', errors='replace')
                            except:
                                row[field] = str(value)
                        else:
                            row[field] = str(value)
                    
                    writer.writerow(row)
                    
            return True
        except Exception as e:
            print(f"导出CSV失败: {str(e)}")
            return False 